package parser

import "vitess.io/vitess/go/vt/sqlparser"

func updateExprsToFields(expr sqlparser.UpdateExprs) ([]*Field, error) {
	fields := make([]*Field, len(expr))
	for i, updateExpr := range expr {
		field := &Field{
			name: updateExpr.Name.Name.String(),
		}
		field.SetExpr(SQLNodeToString(updateExpr.Expr))
		fields[i] = field
	}
	return fields, nil
}

func convertUpdate(u *sqlparser.Update, query string) (Statement, error) {
	tables, err := tableExprsToTables(u.TableExprs)
	if err != nil {
		return nil, err
	}
	stmt := &UpdateStatement{
		Columns: ColumnSet{},
	}
	stmt.sql = query
	stmt.tbl = tables[0].Name
	err = SQLNodeToColumns(stmt.Columns, u)
	if err != nil {
		return nil, err
	}
	fields, err := updateExprsToFields(u.Exprs)
	if err != nil {
		return nil, err
	}
	stmt.Fields = fields
	if u.Where != nil {
		where, err := whereToWhereExpr(u.Where)
		if err != nil {
			return nil, err
		}
		stmt.where = where
	}
	return stmt, nil
}

type UpdateStatement struct {
	BaseStatement
	Fields  []*Field
	Columns ColumnSet
	tbl     string
	where   *WhereExpr
}

func (s *UpdateStatement) Table() string {
	return s.tbl
}

func (s *UpdateStatement) Where() *WhereExpr {
	return s.where
}
